{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/gym/core.py:26: UserWarning: \u001b[33mWARN: Gym minimally supports python 3.6 as the python foundation not longer supports the version, please update your version to 3.7+\u001b[0m\n",
      "  \"Gym minimally supports python 3.6 as the python foundation not longer supports the version, please update your version to 3.7+\"\n",
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/gym/envs/registration.py:593: UserWarning: \u001b[33mWARN: The environment CartPole-v0 is out of date. You should consider upgrading to version `v1`.\u001b[0m\n",
      "  f\"The environment {id} is out of date. You should consider \"\n",
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/gym/core.py:330: DeprecationWarning: \u001b[33mWARN: Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001b[0m\n",
      "  \"Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\"\n",
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/gym/wrappers/step_api_compatibility.py:40: DeprecationWarning: \u001b[33mWARN: Initializing environment in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001b[0m\n",
      "  \"Initializing environment in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\"\n",
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/gym/core.py:52: DeprecationWarning: \u001b[33mWARN: The argument mode in render method is deprecated; use render_mode during environment initialization instead.\n",
      "See here for more information: https://www.gymlibrary.ml/content/api/\u001b[0m\n",
      "  \"The argument mode in render method is deprecated; \"\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAW4AAAD8CAYAAABXe05zAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAUEUlEQVR4nO3db4xd9Z3f8ffHgzGEsIq9DMjYpripIxWSYrYjbyKqiibZ4NLVmjxI5SiN/ADJQSJSoq7Swq7UTR5Y2m42SZ80iZwGxaXZEEsJxY1IuyxNFKEkGEMMwTYGB7xg7GBjlvIvGI/n2wdzHG48/67nj+89M++XdHXP/Z3fuff7s8Yf//ybc+5JVSFJao9FvS5AknR2DG5JahmDW5JaxuCWpJYxuCWpZQxuSWqZOQvuJOuT7E9yIMltc/U5krTQZC7O404yADwJ/BFwCHgI+HhV7Z31D5OkBWauZtzrgANV9XRVvQXcBWyYo8+SpAXlvDl63xXAcx2vDwF/OFHnSy65pK688so5KkWS2ufgwYO8+OKLGW/fXAX3eB/2O2sySTYDmwGuuOIKdu3aNUelSFL7DA0NTbhvrpZKDgGrOl6vBA53dqiqrVU1VFVDg4ODc1SGJM0/cxXcDwFrkqxOcj6wEdgxR58lSQvKnCyVVNVwkk8D/wcYAO6oqj1z8VmStNDM1Ro3VXUvcO9cvb8kLVReOSlJLWNwS1LLGNyS1DIGtyS1jMEtSS1jcEtSyxjcktQyBrcktYzBLUktY3BLUssY3JLUMga3JLWMwS1JLWNwS1LLGNyS1DIGtyS1jMEtSS1jcEtSy8zo1mVJDgKvAqeA4aoaSrIM+C5wJXAQ+LdV9Q8zK1OSdNpszLj/VVWtraqh5vVtwP1VtQa4v3ktSZolc7FUsgHY1mxvA26ag8+QpAVrpsFdwN8meTjJ5qbtsqo6AtA8XzrDz5AkdZjRGjdwXVUdTnIpcF+SJ7o9sAn6zQBXXHHFDMuQpIVjRjPuqjrcPB8F7gbWAS8kWQ7QPB+d4NitVTVUVUODg4MzKUOSFpRpB3eSi5JcfHob+AjwOLAD2NR02wTcM9MiJUlvm8lSyWXA3UlOv8/fVNX/TvIQsD3JzcCzwMdmXqYk6bRpB3dVPQ1cM077ceBDMylKkjQxr5yUpJYxuCWpZQxuSWoZg1uSWsbglqSWMbglqWUMbklqGYNbklrG4JakljG4JallDG5JahmDW5JaxuCWpJYxuCWpZQxuSWoZg1uSWsbglqSWMbglqWUMbklqmSmDO8kdSY4mebyjbVmS+5I81Twv7dh3e5IDSfYnuWGuCpekhaqbGfe3gPVntN0G3F9Va4D7m9ckuQrYCFzdHPPVJAOzVq0kaergrqqfAC+d0bwB2NZsbwNu6mi/q6pOVNUzwAFg3eyUKkmC6a9xX1ZVRwCa50ub9hXAcx39DjVtYyTZnGRXkl3Hjh2bZhmStPDM9i8nM05bjdexqrZW1VBVDQ0ODs5yGZI0f003uF9IshygeT7atB8CVnX0Wwkcnn55kqQzTTe4dwCbmu1NwD0d7RuTLEmyGlgD7JxZiZKkTudN1SHJd4DrgUuSHAL+AvhLYHuSm4FngY8BVNWeJNuBvcAwcGtVnZqj2iVpQZoyuKvq4xPs+tAE/bcAW2ZSlCRpYl45KUktY3BLUssY3JLUMga3JLWMwS1JLWNwS1LLGNyS1DIGtyS1jMEtSS1jcEtSyxjcktQyBrcktYzBLUktY3BLUssY3JLUMga3JLWMwS1JLWNwS1LLTBncSe5IcjTJ4x1tn0/yfJLdzePGjn23JzmQZH+SG+aqcElaqLqZcX8LWD9O+1eqam3zuBcgyVXARuDq5pivJhmYrWIlSV0Ed1X9BHipy/fbANxVVSeq6hngALBuBvVJks4wkzXuTyd5rFlKWdq0rQCe6+hzqGkbI8nmJLuS7Dp27NgMypCkhWW6wf014N3AWuAI8KWmPeP0rfHeoKq2VtVQVQ0NDg5OswxJWnimFdxV9UJVnaqqEeAbvL0ccghY1dF1JXB4ZiVKkjpNK7iTLO94+VHg9BknO4CNSZYkWQ2sAXbOrERJUqfzpuqQ5DvA9cAlSQ4BfwFcn2Qto8sgB4FPAVTVniTbgb3AMHBrVZ2ak8olaYGaMrir6uPjNH9zkv5bgC0zKUqSNDGvnJSkljG4JallDG5JahmDW5JaxuCWpJaZ8qwSaSE58cqLnHj1xTHtS37vEpZcfEkPKpLGMrilRlVxbP8DHHnk3jH7Lh/6E1b88z/uQVXSWC6VSL9V1CmvF1P/M7il06oYOXWy11VIUzK4pUYV1KnhXpchTcngln6rqBGDW/3P4JZOq2LEGbdawOCWfqtcKlErGNxSo5xxqyUMbqlRp07y5v97YUx7Fp3nxTfqKwa31BgZOcXJ1/9hTHsGzmPJxb/fg4qk8Rnc0hSSkAEvMlb/MLilKYVFBrf6yJTBnWRVkh8l2ZdkT5LPNO3LktyX5KnmeWnHMbcnOZBkf5Ib5nIA0pxLyCKDW/2jmxn3MPCnVfVPgfcDtya5CrgNuL+q1gD3N69p9m0ErgbWA19NMjAXxUvnQsClEvWVKYO7qo5U1SPN9qvAPmAFsAHY1nTbBtzUbG8A7qqqE1X1DHAAWDfLdUvnTlwqUX85qzXuJFcC1wIPApdV1REYDXfg0qbbCuC5jsMONW1nvtfmJLuS7Dp27Ng0SpfOEZdK1Ge6Du4k7wS+B3y2ql6ZrOs4bTWmoWprVQ1V1dDg4GC3ZUjnXPzlpPpMV8GdZDGjof3tqvp+0/xCkuXN/uXA0ab9ELCq4/CVwOHZKVeaOzUyMv6OBOIJWOof3ZxVEuCbwL6q+nLHrh3ApmZ7E3BPR/vGJEuSrAbWADtnr2RpbpTfxa2W6Ob/f9cBnwR+mWR30/ZnwF8C25PcDDwLfAygqvYk2Q7sZfSMlFurytuKqO+NnBqmasyqntR3pgzuqnqA8detAT40wTFbgC0zqEs655xxqy1cuJMao7ctc8at/mdwS406NWxuqxUMbqnhjYLVFga31Bi9+41TbvU/g1tqvP7is+Oey33h0svJIv+qqH/40yg1ht94hfFm3Oe/cynxAhz1EX8apSlkYHGvS5B+h8EtTWHRwGImvpRBOvcMbmkKfhe3+o3BLU1h0cBiJ9zqKwa3NIXRGbfJrf5hcEtTWOQvJ9VnDG4JJv1WQINb/cbglhpV499IwYtv1G/8iZQAqqiRib82fvR+IlJ/MLglRmfbIyPDvS5D6orBLcHojPuUN2pSOxjcEqMz7nLGrZbo5mbBq5L8KMm+JHuSfKZp/3yS55Psbh43dhxze5IDSfYnuWEuByDNiinWuKV+0s21vMPAn1bVI0kuBh5Ocl+z7ytV9dednZNcBWwErgYuB/4uyXu8YbD6WVU138ct9b8pZ9xVdaSqHmm2XwX2ASsmOWQDcFdVnaiqZ4ADwLrZKFaaOyPOuNUaZ7XGneRK4Frgwabp00keS3JHkqVN2wrguY7DDjF50Es9V6eGGX7rjbE7sohFi5ec+4KkSXQd3EneCXwP+GxVvQJ8DXg3sBY4AnzpdNdxDh9zWVqSzUl2Jdl17Nixs61bmlXDJ37Dmy//ekz7eRdcxIVLL+9BRdLEugruJIsZDe1vV9X3Aarqhao6VaOXm32Dt5dDDgGrOg5fCRw+8z2ramtVDVXV0ODg4EzGIM2ZZBFZNNDrMqTf0c1ZJQG+Ceyrqi93tC/v6PZR4PFmewewMcmSJKuBNcDO2StZOocSssjv41Z/6eYn8jrgk8Avk+xu2v4M+HiStYwugxwEPgVQVXuSbAf2MnpGyq2eUaK2ShaRAWfc6i9TBndVPcD469b3TnLMFmDLDOqS+kPCImfc6jNeOSlNwhm3+pHBLU3GNW71IYNbAsY5YxWA4Fkl6j8Gt8Qkd8CJN1JQ//EnUgLq1MmJJt1S3zG4JWDk1MlelyB1zeCWgJFhg1vtYXBL4Fe6qlUMbgmXStQuBrdE88tJfzupljC4JZxxq10Mbgl46/WXx20/78KLz20hUhe8llfz1tNPP82vfz325gjjGfj73ePOYl5+63x+9rOfd/Ue73vf+7j4YoNec8/g1rz1xS9+ka9//etd9f2rW/6I69deOab9W//9Trb+r8929R4//elP+cAHPnAWFUrTY3BLjRMjF/D8m+/hzZGLWLb4CJee/yxvnfSr5NV/DG4JODHyDh555SO8PHwpEJ598yre/Y5fcOLkz3pdmjSGv5yUgP2v/yEvD1/G6F+JUAzwqzeu5fAbl/W6NGkMg1sChmsxZ97oqRjgN54lqD7Uzc2CL0iyM8mjSfYk+ULTvizJfUmeap6Xdhxze5IDSfYnuWEuByDNhgsHXuPMC3AG8hY1/EZvCpIm0c2M+wTwwaq6BlgLrE/yfuA24P6qWgPc37wmyVXARuBqYD3w1SR+E7362nve8RCXLznAAKNXUJ6f33D1RQ9wEc/3ujRpjG5uFlzAa83Lxc2jgA3A9U37NuDHwH9s2u+qqhPAM0kOAOuACX/Lc/Lkya7Pt5W69cYb3c+W/+dPHuWyPX/P8ZMreGvkQn7vvBd5cOA4Tzz7Ytfv8dJLL/lzrFlz8uTE63RdnVXSzJgfBv4J8F+r6sEkl1XVEYCqOpLk0qb7CqDzioVDTduEjh8/zp133tlNKVLXnnzyya77PrjveeB5YM+0P++HP/whe/funfbxUqfjx49PuK+r4K6qU8DaJO8C7k7y3km6Z5y2Md/ek2QzsBngiiuu4HOf+1w3pUhde/rpp/n5z7u76nE2fOITn/ACHM2a7373uxPuO6uzSqrqZUaXRNYDLyRZDtA8H226HQJWdRy2Ejg8znttraqhqhoaHBw8mzIkaUHr5qySwWamTZILgQ8DTwA7gE1Nt03APc32DmBjkiVJVgNrgJ2zXLckLVjdLJUsB7Y169yLgO1V9YMkPwO2J7kZeBb4GEBV7UmyHdgLDAO3NkstkqRZ0M1ZJY8B147Tfhz40ATHbAG2zLg6SdIYXjkpSS1jcEtSy/jtgJq3rrnmGm666aZz9nnLli07Z5+lhc3g1rx1yy23cMstt/S6DGnWuVQiSS1jcEtSyxjcktQyBrcktYzBLUktY3BLUssY3JLUMga3JLWMwS1JLWNwS1LLGNyS1DIGtyS1jMEtSS1jcEtSy3Rzs+ALkuxM8miSPUm+0LR/PsnzSXY3jxs7jrk9yYEk+5PcMJcDkKSFppvv4z4BfLCqXkuyGHggyQ+bfV+pqr/u7JzkKmAjcDVwOfB3Sd7jDYMlaXZMOeOuUa81Lxc3j5rkkA3AXVV1oqqeAQ4A62ZcqSQJ6HKNO8lAkt3AUeC+qnqw2fXpJI8luSPJ0qZtBfBcx+GHmjZJ0izoKrir6lRVrQVWAuuSvBf4GvBuYC1wBPhS0z3jvcWZDUk2J9mVZNexY8emUbokLUxndVZJVb0M/BhYX1UvNIE+AnyDt5dDDgGrOg5bCRwe5722VtVQVQ0NDg5Op3ZJWpC6OatkMMm7mu0LgQ8DTyRZ3tHto8DjzfYOYGOSJUlWA2uAnbNatSQtYN2cVbIc2JZkgNGg315VP0hyZ5K1jC6DHAQ+BVBVe5JsB/YCw8CtnlEiSbNnyuCuqseAa8dp/+Qkx2wBtsysNEnSeLxyUpJaxuCWpJYxuCWpZQxuSWoZg1uSWsbglqSWMbglqWUMbklqGYNbklrG4JakljG4JallDG5JahmDW5JaxuCWpJYxuCWpZQxuSWoZg1uSWsbglqSWMbglqWUMbklqGYNbklrG4JaklklV9boGkhwDXgde7HUtc+ASHFfbzNexOa52+UdVNTjejr4IboAku6pqqNd1zDbH1T7zdWyOa/5wqUSSWsbglqSW6afg3trrAuaI42qf+To2xzVP9M0atySpO/0045YkdaHnwZ1kfZL9SQ4kua3X9ZytJHckOZrk8Y62ZUnuS/JU87y0Y9/tzVj3J7mhN1VPLcmqJD9Ksi/JniSfadpbPbYkFyTZmeTRZlxfaNpbPa7Tkgwk+UWSHzSv58u4Dib5ZZLdSXY1bfNibNNSVT17AAPAr4B/DJwPPApc1cuapjGGfwn8AfB4R9tfAbc127cB/7nZvqoZ4xJgdTP2gV6PYYJxLQf+oNm+GHiyqb/VYwMCvLPZXgw8CLy/7ePqGN+/B/4G+MF8+Vls6j0IXHJG27wY23QevZ5xrwMOVNXTVfUWcBewocc1nZWq+gnw0hnNG4BtzfY24KaO9ruq6kRVPQMcYPTPoO9U1ZGqeqTZfhXYB6yg5WOrUa81Lxc3j6Ll4wJIshL4N8B/62hu/bgmMZ/HNqleB/cK4LmO14eatra7rKqOwGgAApc27a0cb5IrgWsZnZ22fmzNcsJu4ChwX1XNi3EB/wX4D8BIR9t8GBeM/uP6t0keTrK5aZsvYztr5/X48zNO23w+zaV1403yTuB7wGer6pVkvCGMdh2nrS/HVlWngLVJ3gXcneS9k3RvxbiS/DFwtKoeTnJ9N4eM09Z34+pwXVUdTnIpcF+SJybp27axnbVez7gPAas6Xq8EDveoltn0QpLlAM3z0aa9VeNNspjR0P52VX2/aZ4XYwOoqpeBHwPraf+4rgP+JMlBRpccP5jkf9D+cQFQVYeb56PA3YwufcyLsU1Hr4P7IWBNktVJzgc2Ajt6XNNs2AFsarY3Afd0tG9MsiTJamANsLMH9U0po1PrbwL7qurLHbtaPbYkg81MmyQXAh8GnqDl46qq26tqZVVdyejfo/9bVf+Olo8LIMlFSS4+vQ18BHiceTC2aev1b0eBGxk9Y+FXwJ/3up5p1P8d4AhwktF/6W8Gfh+4H3iqeV7W0f/Pm7HuB/51r+ufZFz/gtH/Xj4G7G4eN7Z9bMA/A37RjOtx4D817a0e1xljvJ63zypp/bgYPevs0eax53ROzIexTffhlZOS1DK9XiqRJJ0lg1uSWsbglqSWMbglqWUMbklqGYNbklrG4JakljG4Jall/j8IwByH/8Tu8QAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "#创建环境\n",
    "env = gym.make('CartPole-v0')\n",
    "env.reset()\n",
    "\n",
    "\n",
    "#打印游戏\n",
    "def show():\n",
    "    plt.imshow(env.render(mode='rgb_array'))\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, tensor(0.3800, grad_fn=<SelectBackward0>))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import random\n",
    "\n",
    "\n",
    "#定义模型\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.sequential = torch.nn.Sequential(\n",
    "            torch.nn.Linear(4, 128),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(128, 2),\n",
    "            torch.nn.Softmax(dim=1),\n",
    "        )\n",
    "\n",
    "    #得到一个动作\n",
    "    def forward(self, state):\n",
    "        #[1, 4] -> [1, 2]\n",
    "        prob = self.sequential(state)\n",
    "\n",
    "        #根据概率选择一个动作\n",
    "        action = random.choices(range(2), weights=prob[0].tolist(), k=1)[0]\n",
    "\n",
    "        return action, prob[0, action]\n",
    "\n",
    "\n",
    "model = Model()\n",
    "\n",
    "model(torch.randn(1, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2994]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_td = sequential = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 128),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(128, 1),\n",
    ")\n",
    "\n",
    "model_td(torch.randn(1, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([tensor([[ 0.0103, -0.0136, -0.0174,  0.0115]]),\n",
       "  tensor([[ 0.0100, -0.2084, -0.0172,  0.2986]]),\n",
       "  tensor([[ 0.0058, -0.4033, -0.0112,  0.5858]]),\n",
       "  tensor([[-2.2255e-03, -5.9827e-01,  4.8581e-04,  8.7493e-01]]),\n",
       "  tensor([[-0.0142, -0.7934,  0.0180,  1.1678]]),\n",
       "  tensor([[-0.0301, -0.9887,  0.0413,  1.4660]]),\n",
       "  tensor([[-0.0498, -1.1844,  0.0707,  1.7713]]),\n",
       "  tensor([[-0.0735, -0.9901,  0.1061,  1.5014]]),\n",
       "  tensor([[-0.0933, -1.1863,  0.1361,  1.8253]]),\n",
       "  tensor([[-0.1170, -0.9930,  0.1726,  1.5778]]),\n",
       "  tensor([[-0.1369, -0.8003,  0.2042,  1.3435]])],\n",
       " [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],\n",
       " [tensor(0.5552, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.5612, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.5742, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.5999, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.6297, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.6574, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.3166, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.6603, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.3130, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.3337, grad_fn=<SelectBackward0>),\n",
       "  tensor(0.6449, grad_fn=<SelectBackward0>)],\n",
       " [tensor([[ 0.0100, -0.2084, -0.0172,  0.2986]]),\n",
       "  tensor([[ 0.0058, -0.4033, -0.0112,  0.5858]]),\n",
       "  tensor([[-2.2255e-03, -5.9827e-01,  4.8581e-04,  8.7493e-01]]),\n",
       "  tensor([[-0.0142, -0.7934,  0.0180,  1.1678]]),\n",
       "  tensor([[-0.0301, -0.9887,  0.0413,  1.4660]]),\n",
       "  tensor([[-0.0498, -1.1844,  0.0707,  1.7713]]),\n",
       "  tensor([[-0.0735, -0.9901,  0.1061,  1.5014]]),\n",
       "  tensor([[-0.0933, -1.1863,  0.1361,  1.8253]]),\n",
       "  tensor([[-0.1170, -0.9930,  0.1726,  1.5778]]),\n",
       "  tensor([[-0.1369, -0.8003,  0.2042,  1.3435]]),\n",
       "  tensor([[-0.1529, -0.9973,  0.2310,  1.6925]])],\n",
       " [False, False, False, False, False, False, False, False, False, False, True])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_data():\n",
    "    states = []\n",
    "    rewards = []\n",
    "    probs = []\n",
    "    next_states = []\n",
    "    overs = []\n",
    "\n",
    "    #初始化游戏\n",
    "    state = env.reset()\n",
    "\n",
    "    #玩到游戏结束为止\n",
    "    over = False\n",
    "    while not over:\n",
    "        #根据当前状态得到一个动作\n",
    "        state = torch.FloatTensor(state).reshape(1, 4)\n",
    "        action, prob = model(state)\n",
    "\n",
    "        #执行动作,得到反馈\n",
    "        next_state, reward, over, _ = env.step(action)\n",
    "\n",
    "        #记录数据样本\n",
    "        states.append(state)\n",
    "        rewards.append(reward)\n",
    "        probs.append(prob)\n",
    "        next_states.append(torch.FloatTensor(next_state).reshape(1, 4))\n",
    "        overs.append(over)\n",
    "\n",
    "        #更新游戏状态,开始下一个动作\n",
    "        state = next_state\n",
    "\n",
    "    return states, rewards, probs, next_states, overs\n",
    "\n",
    "\n",
    "get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "17.0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "\n",
    "\n",
    "def test(play):\n",
    "    #初始化游戏\n",
    "    state = env.reset()\n",
    "\n",
    "    #记录反馈值的和,这个值越大越好\n",
    "    reward_sum = 0\n",
    "\n",
    "    #玩到游戏结束为止\n",
    "    over = False\n",
    "    while not over:\n",
    "        #根据当前状态得到一个动作\n",
    "        state = torch.FloatTensor(state).reshape(1, 4)\n",
    "        with torch.no_grad():\n",
    "            action, _ = model(state)\n",
    "\n",
    "        #执行动作,得到反馈\n",
    "        state, reward, over, _ = env.step(action)\n",
    "        reward_sum += reward\n",
    "\n",
    "        #打印动画\n",
    "        if play:\n",
    "            display.clear_output(wait=True)\n",
    "            show()\n",
    "\n",
    "    return reward_sum\n",
    "\n",
    "\n",
    "test(play=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "executionInfo": {
     "elapsed": 8251,
     "status": "ok",
     "timestamp": 1650011468229,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "BQXVYW2T_DcQ",
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 26.7\n",
      "100 15.8\n",
      "200 43.0\n",
      "300 88.1\n",
      "400 78.4\n",
      "500 183.3\n",
      "600 182.7\n",
      "700 194.5\n",
      "800 200.0\n",
      "900 188.7\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "    optimizer_td = torch.optim.Adam(model_td.parameters(), lr=1e-2)\n",
    "\n",
    "    #玩N局游戏,每局游戏训练一次\n",
    "    for epoch in range(1000):\n",
    "        #玩一局游戏,得到数据\n",
    "        states, rewards, probs, next_states, overs = get_data()\n",
    "\n",
    "        #计算values和targets\n",
    "        values = [model_td(state) for state in states]\n",
    "        targets = [model_td(next_state) * 0.98 for next_state in next_states]\n",
    "        targets = [\n",
    "            target * (0 if over else 1)\n",
    "            for target, over in zip(targets, overs)\n",
    "        ]\n",
    "        targets = [target + reward for target, reward in zip(targets, rewards)]\n",
    "\n",
    "        #取出每一步的动作概率,取对数后作为loss\n",
    "        #因为概率一定是0-1之间的小数,所以取对数之后一定是负数\n",
    "        #符号取反,一定是正数,最终导致的结果就是,概率越小,loss越大,反之亦然\n",
    "        #这样会迫使prob的取值越来越大\n",
    "        probs = [-i.log() for i in probs]\n",
    "\n",
    "        #第一部分的loss是针对prob的,以target和value的差作为系数\n",
    "        loss = [\n",
    "            prob * (target - value).detach()\n",
    "            for prob, target, value in zip(probs, targets, values)\n",
    "        ]\n",
    "        loss = sum(loss) / len(loss)\n",
    "\n",
    "        #第二部分的loss是针对value的,越接近target就越好\n",
    "        loss_td = [\n",
    "            torch.nn.functional.mse_loss(value, target.detach())\n",
    "            for value, target in zip(values, targets)\n",
    "        ]\n",
    "        loss_td = sum(loss_td) / len(loss_td)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        optimizer_td.zero_grad()\n",
    "        loss.backward()\n",
    "        loss_td.backward()\n",
    "        optimizer.step()\n",
    "        optimizer_td.step()\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            test_result = sum([test(play=False) for _ in range(10)]) / 10\n",
    "            print(epoch, test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAW4AAAD8CAYAAABXe05zAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAWA0lEQVR4nO3daYxd5X3H8e9vFo893rDxGBwv2CGmrUmKQSOXiiiiSRq7NKqJVCqjNrgSlfMCpESN0kKiNskLS2mVpW9KJKcgLJrEWEoQLtA2jpsojZRixsSAFwwDdu1hjGdYvNuz3X9f3DPh1j7juXMXZp6Z30e6uuc+Z/s/UfhxeOY55ygiMDOzdDSMdwFmZjY2Dm4zs8Q4uM3MEuPgNjNLjIPbzCwxDm4zs8TULbglrZN0SFKnpAfqdR4zs6lG9ZjHLakReAX4Q6ALeA64OyIO1PxkZmZTTL2uuNcAnRHxekT0A9uA9XU6l5nZlNJUp+MuBo6V/O4Cfm+kjRcsWBDLly+vUylmZuk5cuQIb731lvLW1Su48072/8ZkJG0CNgEsW7aMjo6OOpViZpae9vb2EdfVa6ikC1ha8nsJ0F26QURsiYj2iGhva2urUxlmZpNPvYL7OWClpBWSpgEbgB11OpeZ2ZRSl6GSiBiUdD/wn0Aj8EhE7K/HuczMppp6jXETEc8Az9Tr+GZmU5XvnDQzS4yD28wsMQ5uM7PEOLjNzBLj4DYzS4yD28wsMQ5uM7PEOLjNzBLj4DYzS4yD28wsMQ5uM7PEOLjNzBLj4DYzS4yD28wsMQ5uM7PEOLjNzBLj4DYzS4yD28wsMVW9ukzSEeAMMAQMRkS7pPnA48By4AjwZxHxbnVlmpnZsFpccf9BRKyOiPbs9wPArohYCezKfpuZWY3UY6hkPbA1W94K3FmHc5iZTVnVBncAP5G0R9KmrO2aiDgOkH0vrPIcZmZWoqoxbuC2iOiWtBDYKenlcnfMgn4TwLJly6osw8xs6qjqijsiurPvHuAJYA1wQtIigOy7Z4R9t0REe0S0t7W1VVOGmdmUUnFwS5opafbwMvApYB+wA9iYbbYReLLaIs3M7D3VDJVcAzwhafg4P4iI/5D0HLBd0r3AUeCu6ss0M7NhFQd3RLwO3JTT/jbwiWqKMjOzkfnOSTOzxDi4zcwS4+A2M0uMg9vMLDEObjOzxDi4zcwS4+A2M0uMg9vMLDEObjOzxDi4zcwS4+A2M0uMg9vMLDEObjOzxDi4zcwS4+A2M0uMg9vMLDEObjOzxDi4zcwS4+A2M0vMqMEt6RFJPZL2lbTNl7RT0qvZ97ySdQ9K6pR0SNLaehVuZjZVlXPF/Siw7pK2B4BdEbES2JX9RtIqYANwY7bPQ5Iaa1atmZmNHtwR8QvgnUua1wNbs+WtwJ0l7dsioi8iDgOdwJralGpmZlD5GPc1EXEcIPtemLUvBo6VbNeVtV1G0iZJHZI6ent7KyzDzGzqqfUfJ5XTFnkbRsSWiGiPiPa2trYal2FmNnlVGtwnJC0CyL57svYuYGnJdkuA7srLMzOzS1Ua3DuAjdnyRuDJkvYNklokrQBWArurK9HMzEo1jbaBpB8CtwMLJHUBXwW+AWyXdC9wFLgLICL2S9oOHAAGgfsiYqhOtZuZTUmjBndE3D3Cqk+MsP1mYHM1RZmZ2ch856SZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiRk1uCU9IqlH0r6Stq9JekPS3uxzR8m6ByV1SjokaW29Cjczm6rKueJ+FFiX0/6diFidfZ4BkLQK2ADcmO3zkKTGWhVrZmZlBHdE/AJ4p8zjrQe2RURfRBwGOoE1VdRnZmaXqGaM+35JL2ZDKfOytsXAsZJturK2y0jaJKlDUkdvb28VZZiZTS2VBvd3geuB1cBx4FtZu3K2jbwDRMSWiGiPiPa2trYKyzAzm3oqCu6IOBERQxFRAL7He8MhXcDSkk2XAN3VlWhmZqUqCm5Ji0p+fgYYnnGyA9ggqUXSCmAlsLu6Es3MrFTTaBtI+iFwO7BAUhfwVeB2SaspDoMcAT4HEBH7JW0HDgCDwH0RMVSXys3MpqhRgzsi7s5pfvgK228GNldTlJmZjcx3TpqZJcbBbWaWGAe3mVliHNxmZolxcJuZJWbUWSVmZla9i6d7KfRfpLl1Dk3TZ6OGyq+bHdxmZnUWhSHefOEnvPv6Hppb5zJt1nxaFyxj5oJlzLr2eppnzBnT8RzcZmZ1VhgaoO/MWwxePMvgxbNceOcNTh19CYAlt/4pi2761JiO5zFuM7M6G+w7z5nuV3LXSXnP5rsyB7eZWb0FEJc/KLVx2gxmLvzgmA/n4DYzq7PCUD95T7huaGymZdb8MR/PwW1mVmf9Z/JfIqbGJpqmzxzz8RzcZmZ1dqprP1EY4UGpHuM2M5t4Bi+cyW1vXbAUaewx7OA2M6ujiMKIV9szF1wHDm4zs4mlMDjAYN+53HXNM6+q6JgObjOzOhrsO0ff6d7cdWpo8jxuM7OJZmiE4FZjE43TZlR0zFGDW9JSST+TdFDSfkmfz9rnS9op6dXse17JPg9K6pR0SNLaiiozM5sEIufGG4Dm1quYuWBZRccs54p7EPhiRPwOcCtwn6RVwAPArohYCezKfpOt2wDcCKwDHpLUWFF1ZmaJGzh/Ore9sWkaTdNnVXTMUYM7Io5HxPPZ8hngILAYWA9szTbbCtyZLa8HtkVEX0QcBjqBNRVVZ2aWuL7TPfkrGhpRY2XP+RvTGLek5cDNwLPANRFxHIrhDizMNlsMHCvZrStru/RYmyR1SOro7c0fuDczS92pYwdy2yv5o+SwsoNb0izgR8AXIiL/2j/bNKftskGeiNgSEe0R0d7W1lZuGWZmyYiIEedwX3XdTRUft6zgltRMMbS/HxE/zppPSFqUrV8EDP/3QBewtGT3JUB3xRWamSUqhgYoDPbnrpsxb1HFxy1nVomAh4GDEfHtklU7gI3Z8kbgyZL2DZJaJK0AVgK7K67QzCxRAxfOMHDhVO66xmkzKh4uKWdk/Dbgs8BLkvZmbV8GvgFsl3QvcBS4CyAi9kvaDhygOCPlvogY4ekqZmaTV9/pHvpO5fxxUg2oofLJdqMGd0T8kvxxa4BPjLDPZmBzxVWZmSVupPnbANPntNEyd+GI60fjOyfNzOpkaKAvt72xpZWmltaKj+vgNjOrk74zb+e2NzS10NDUUvFxHdxmZnVy8sje/BV6n+Zxm5nZ2Az1X8htn33th6o6roPbzKwOojA04h8oW69emtteLge3mVkdDPWdozCY/8fJ5ta5VR3bwW1mVgcXT/WM+K7JauZwg4PbzKwu+s+dzB3jbmyZWdWMEnBwm5nVXERAFHLXtc5fzLSZHioxM5tw+s/nP6OkabqvuM3MJp4ILp58M3eVGpqgijnc4OA2M6u5iAJnug/lrlNDY1U334CD28ysPvLmcKuBq1bcXPWhHdxmZjU21H+eQs6bb0TxyYDVcnCbmdXYhXffpDBw8fIVEo3N1f1hEhzcZmY1d/7tY7lzuBuaWqDKm2/AwW1mVlPF55PkP6Nk5sIVNE2r/DncwxzcZmY1FQz15wyTANNa56Km5qrPUM7LgpdK+pmkg5L2S/p81v41SW9I2pt97ijZ50FJnZIOSVpbdZVmZomIQoH+EV6gUHxBcPXXy+W8LHgQ+GJEPC9pNrBH0s5s3Xci4pulG0taBWwAbgQ+APxU0g1+YbCZTQUxNMCpY/vzV1b5AoVho0Z/RByPiOez5TPAQWDxFXZZD2yLiL6IOAx0AmuqrtTMLAERQWFo4LJ2NTYxa9HKmpxjTNfskpYDNwPPZk33S3pR0iOS5mVti4FjJbt1ceWgNzObNIqhffkfJ6VGZsy9tibnKDu4Jc0CfgR8ISJOA98FrgdWA8eBbw1vmrP7Zb2QtElSh6SO3t7esdZtZjYhDZw7SRQufzKgGhpomjG7JucoK7glNVMM7e9HxI8BIuJERAxFRAH4Hu8Nh3QBpe/lWQJ0X3rMiNgSEe0R0d7WVv2dRGZmE8HZE69RGOzPXVeL8W0ob1aJgIeBgxHx7ZL2RSWbfQbYly3vADZIapG0AlgJ7K5JtWZmE1z/2Xdzn8XdMqet+GTAGijnKLcBnwVekrQ3a/sycLek1RSHQY4AnwOIiP2StgMHKM5Iuc8zSsxsKogIYoQXKMxc+EEaajCHG8oI7oj4Jfnj1s9cYZ/NwOYq6jIzS04UBhm8cDp33bSZc6EGc7jBd06amdVMYXCAi6d6ctepoen9G+M2M7PyDA1c5Pzbxy5fIdHQNK1m53Fwm5nVSt7LE4CmlpnMWfzbNTuNg9vMrEYGLpzODW81NtHcWt2b3Us5uM3MauTcicO5s0qEajajBBzcZmY1c673SO4Vd/GOydr8YRIc3GZmNREjjG8DzP7Ab9Xkca7DHNxmZjUQhSGG8t4zCbTMuhpqNBUQHNxmZjVRGLjIwPlTuesaW2bU9FwObjOzGug/f4pzvf+bu05qqNnNN+DgNjOrjYjch0s1t86l9eqlOTtUzsFtZlYDI41vNzRPp3lm7eZwQ3lPBzQzm5I6Ozvp6cl/9silpp09mtt+sa+f3Xv2Mtp0QEl85CMfYdasWaOey8FtZjaCzZs38+ijj5a17d/d8zE+/fs3XDaW3dnZyZ//5UfLOsaePXu45ZZbRt3OwW1mVgNXz2mlP1rpunADfYVWrm7upm3aUV4++lbNz+XgNjOrUoPEALPZc3otpwYXAOLoxVWsbO3gxdf/u/bnq/kRzcymmJnTmzlW+BinBtsoxqoIGnn1fDuv9M6p+fkc3GZmVZo/Zwbz5szl0j9ABo0MFGo3f3tYOS8Lni5pt6QXJO2X9PWsfb6knZJezb7nlezzoKROSYckra151WZmE8iCua1ctyAovoL3PYMD5xkcOF/z85Vzxd0HfDwibgJWA+sk3Qo8AOyKiJXAruw3klYBG4AbgXXAQ5Iaa165mdkEcbF/kNkXfkrr4AEUA0AwTeeZc+7feLf3UM3PV87LggM4m/1szj4BrAduz9q3Aj8H/jZr3xYRfcBhSZ3AGuBXI51jYGCAN998s7IemJnVyYULF8rabv+RXv7qG9tZsnAXs6++kWvbrqF9eYHCwAnOXOgv+3xvv/32b7JwYGBgxO3KmlWSXTHvAT4E/HNEPCvpmog4DhARxyUtzDZfDPxPye5dWdsVi33sscfKKcXM7H3z2muvlb3tqXN9nDrcDYe7aZB4ulEIUSiM/LjXSz399NPs3bsXKObiSMoK7ogYAlZLugp4QtKHr7B53kj8ZZVL2gRsAli2bBlf+tKXyinFzOx9c+DAATo6Osa8XyGCwmD5gT3snnvu+c0NOI8//viI241pVklEnKQ4JLIOOCFpEUD2PXxfaBdQ+kSVJUB3zrG2RER7RLS3tbWNpQwzsymtnFklbdmVNpJmAJ8EXgZ2ABuzzTYCT2bLO4ANklokrQBWArtrXLeZ2ZRVzlDJImBrNs7dAGyPiKck/QrYLule4ChwF0BE7Je0HTgADAL3ZUMtZmZWA+XMKnkRuDmn/W3gEyPssxnYXHV1ZmZ2Gd85aWaWGAe3mVli/HRAM7MR3HLLLZw8efJ9O9/cueW9KUfFGyPHV3t7e1QyV9LMbLJqb2+no6Mj9wlVHioxM0uMg9vMLDEObjOzxDi4zcwS4+A2M0uMg9vMLDEObjOzxDi4zcwS4+A2M0uMg9vMLDEObjOzxDi4zcwS4+A2M0uMg9vMLDHlvCx4uqTdkl6QtF/S17P2r0l6Q9Le7HNHyT4PSuqUdEjS2np2wMxsqinnRQp9wMcj4qykZuCXkv49W/ediPhm6caSVgEbgBuBDwA/lXSDXxhsZlYbo15xR9HZ7Gdz9rnS2xfWA9sioi8iDgOdwJqqKzUzM6DMMW5JjZL2Aj3Azoh4Nlt1v6QXJT0iaV7Wthg4VrJ7V9ZmZmY1UFZwR8RQRKwGlgBrJH0Y+C5wPbAaOA58K9s871U7l12hS9okqUNSR29vbwWlm5lNTWOaVRIRJ4GfA+si4kQW6AXge7w3HNIFLC3ZbQnQnXOsLRHRHhHtbW1tldRuZjYllTOrpE3SVdnyDOCTwMuSFpVs9hlgX7a8A9ggqUXSCmAlsLumVZuZTWHlzCpZBGyV1Egx6LdHxFOSHpO0muIwyBHgcwARsV/SduAAMAjc5xklZma1M2pwR8SLwM057Z+9wj6bgc3VlWZmZnl856SZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiXFwm5klxsFtZpYYB7eZWWIc3GZmiVFEjHcNSOoFzgFvjXctdbAA9ys1k7Vv7ldarouItrwVEyK4ASR1RET7eNdRa+5XeiZr39yvycNDJWZmiXFwm5klZiIF95bxLqBO3K/0TNa+uV+TxIQZ4zYzs/JMpCtuMzMrw7gHt6R1kg5J6pT0wHjXM1aSHpHUI2lfSdt8STslvZp9zytZ92DW10OS1o5P1aOTtFTSzyQdlLRf0uez9qT7Jmm6pN2SXsj69fWsPel+DZPUKOnXkp7Kfk+Wfh2R9JKkvZI6srZJ0beKRMS4fYBG4DXgg8A04AVg1XjWVEEfPgbcAuwraftH4IFs+QHgH7LlVVkfW4AVWd8bx7sPI/RrEXBLtjwbeCWrP+m+AQJmZcvNwLPAran3q6R/fw38AHhqsvx/Mav3CLDgkrZJ0bdKPuN9xb0G6IyI1yOiH9gGrB/nmsYkIn4BvHNJ83pga7a8FbizpH1bRPRFxGGgk+L/BhNORByPiOez5TPAQWAxifctis5mP5uzT5B4vwAkLQH+GPiXkubk+3UFk7lvVzTewb0YOFbyuytrS901EXEcigEILMzak+yvpOXAzRSvTpPvWzacsBfoAXZGxKToF/BPwN8AhZK2ydAvKP7L9SeS9kjalLVNlr6NWdM4n185bZN5mkty/ZU0C/gR8IWIOC3ldaG4aU7bhOxbRAwBqyVdBTwh6cNX2DyJfkn6NNATEXsk3V7OLjltE65fJW6LiG5JC4Gdkl6+wrap9W3MxvuKuwtYWvJ7CdA9TrXU0glJiwCy756sPan+SmqmGNrfj4gfZ82Tom8AEXES+DmwjvT7dRvwJ5KOUBxy/LikfyX9fgEQEd3Zdw/wBMWhj0nRt0qMd3A/B6yUtELSNGADsGOca6qFHcDGbHkj8GRJ+wZJLZJWACuB3eNQ36hUvLR+GDgYEd8uWZV03yS1ZVfaSJoBfBJ4mcT7FREPRsSSiFhO8Z+j/4qIvyDxfgFImilp9vAy8ClgH5OgbxUb77+OAndQnLHwGvCV8a6ngvp/CBwHBij+m/5e4GpgF/Bq9j2/ZPuvZH09BPzReNd/hX59lOJ/Xr4I7M0+d6TeN+B3gV9n/doH/H3WnnS/Lunj7bw3qyT5flGcdfZC9tk/nBOToW+VfnznpJlZYsZ7qMTMzMbIwW1mlhgHt5lZYhzcZmaJcXCbmSXGwW1mlhgHt5lZYhzcZmaJ+T9CQbsh2GtU/AAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test(play=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "第9章-策略梯度算法.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
